import pymysql
import logging

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 数据库连接配置
DB_CONFIG = {
    'host': '150.109.6.243',
    'user': 'music_ocr_db',
    'password': 'wZktkPy4X2TpZcCY',
    'db': 'music_ocr_db',
    'charset': 'utf8mb4',
}

def check_and_fix_chat_messages_table():
    """检查并修复chat_messages表结构"""
    connection = pymysql.connect(**DB_CONFIG)
    try:
        with connection.cursor() as cursor:
            # 检查表是否存在
            cursor.execute("SHOW TABLES LIKE 'chat_messages'")
            if not cursor.fetchone():
                logger.info("chat_messages表不存在，将创建新表")
                create_chat_messages_table(cursor)
                connection.commit()
                return True
            
            # 检查表结构
            cursor.execute("SHOW COLUMNS FROM chat_messages")
            columns = [column[0] for column in cursor.fetchall()]  # 使用列索引而不是字典键
            
            # 检查是否缺少必要的列
            required_columns = ['conversation_id', 'role', 'content', 'user_id', 'created_at']
            missing_columns = [col for col in required_columns if col not in columns]
            
            if missing_columns:
                logger.info(f"chat_messages表缺少以下列: {missing_columns}")
                # 删除旧表并创建新表
                logger.info("将删除旧表并创建新表")
                cursor.execute("DROP TABLE IF EXISTS chat_messages")
                create_chat_messages_table(cursor)
                connection.commit()
                return True
            else:
                logger.info("chat_messages表结构正确")
                return True
    except Exception as e:
        logger.error(f"检查表结构时出错: {str(e)}")
        return False
    finally:
        connection.close()

def create_chat_messages_table(cursor):
    """创建chat_messages表"""
    sql = """
    CREATE TABLE chat_messages (
        id INT AUTO_INCREMENT PRIMARY KEY,
        conversation_id INT NOT NULL,
        user_id INT NOT NULL,
        role ENUM('user', 'assistant', 'system') NOT NULL,
        content TEXT NOT NULL,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE,
        FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
    """
    cursor.execute(sql)
    # 创建索引
    cursor.execute("CREATE INDEX idx_chat_messages_conversation ON chat_messages(conversation_id)")
    cursor.execute("CREATE INDEX idx_chat_messages_user ON chat_messages(user_id)")
    cursor.execute("CREATE INDEX idx_chat_messages_created ON chat_messages(created_at)")
    logger.info("chat_messages表创建成功")

if __name__ == "__main__":
    logger.info("开始检查并修复chat_messages表...")
    if check_and_fix_chat_messages_table():
        logger.info("表结构检查和修复完成")
    else:
        logger.error("表结构检查和修复失败") 